#include "ESVideo.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#ifdef __cplusplus
extern "C" {
#endif

#include "libavcodec/avcodec.h"

#ifdef __cplusplus
}
#endif

typedef struct H264Dec
{
    AVCodecContext *c;
    AVFrame *picture;
}H264Dec;

void* H264dec_open(void)
{
    H264Dec * pH264Dec = new H264Dec;
    AVCodec *codec;

    pH264Dec->c = NULL;
    avcodec_register_all();

    /* find the mpeg1 video decoder */
    codec = avcodec_find_decoder(AV_CODEC_ID_H264);

    if (!codec) {
        fprintf(stderr, "codec not found\n");
        return 0;
    }

    pH264Dec->c= avcodec_alloc_context3(codec);
    pH264Dec->picture= avcodec_alloc_frame();

    if(codec->capabilities&CODEC_CAP_TRUNCATED)
        pH264Dec->c->flags|= 0;

    /* for some codecs, such as msmpeg4 and mpeg4, width and height
       MUST be initialized there because these info are not available
       in the bitstream */

    /* open it */
    AVDictionary* opts;
    av_dict_set(&opts, "b", "2.5M", 0);
    if (avcodec_open2(pH264Dec->c, codec, &opts) < 0) {
        fprintf(stderr, "could not open codec\n");
       return 0;
    }

    return pH264Dec;
}

void H264dec_close(void *pH264dec)
{
    H264Dec * pH264DecTemp = (H264Dec *)pH264dec;

    avcodec_close( (AVCodecContext *) (pH264DecTemp->c));
    av_free( (AVCodecContext *)  pH264DecTemp->c);
    av_free( (AVFrame *) pH264DecTemp->picture);

    free(pH264DecTemp);

    return;
}

int H264dec_decode(void *pH264dec, int *pGotPic, unsigned char *pBsBuf, int nBsSize, H264Frame *frame)
{
    int i;
    int ret;
    H264Dec * pH264DecTemp = (H264Dec *)pH264dec;

    try
    {
        AVPacket packet;
        av_init_packet(&packet);
        packet.data = (uint8_t*)pBsBuf;
        packet.size = nBsSize;
        ret = avcodec_decode_video2( (AVCodecContext *) pH264DecTemp->c,  (AVFrame *) pH264DecTemp->picture, \
                                     pGotPic, &packet);
        av_destruct_packet(&packet);
    }
    catch(...)
    {
        *pGotPic = 0;
        return 0;
    }

    for (i = 0; i < 3; i++)
    {
        frame->pitch[i] = pH264DecTemp->picture->linesize[i];
    }

    //yuv out put
    if (*pGotPic && frame->yuv[0]!=NULL)
    {
        int i;
        int xsize, ysize, wrap;
        unsigned char *dst;
        unsigned char *src;

        {
            unsigned long yuv_size = pH264DecTemp->c->width*pH264DecTemp->c->height;
            if (frame->width*frame->height < yuv_size)
            {

                frame->width = pH264DecTemp->c->width;
                frame->height = pH264DecTemp->c->height;
                frame->key_frame = pH264DecTemp->picture->key_frame;	//get key frame flag

                return ret;

            }
            else
            {
                frame->width = pH264DecTemp->c->width;
                frame->height = pH264DecTemp->c->height;
                frame->key_frame = pH264DecTemp->picture->key_frame;	//get key frame flag
            }
        }

        src = pH264DecTemp->picture->data[0];
        dst = frame->yuv[0];
        xsize = frame->width;
        ysize = frame->height;
        wrap = frame->pitch[0];
        for(i=0;i<ysize;i++)
        {
            memcpy(dst,src,xsize);
            src += wrap;
            dst += xsize;
        }

        src = pH264DecTemp->picture->data[1];
        dst = frame->yuv[1];
        xsize =xsize>>1;
        ysize = ysize>>1;
        wrap = frame->pitch[1];
        for(i=0;i< ysize;i++)
        {
            memcpy(dst,src,xsize);
            src += wrap;
            dst += xsize;
        }

        src = pH264DecTemp->picture->data[2];
        dst = frame->yuv[2];
        wrap = frame->pitch[2];
        for(i=0;i< ysize;i++)
        {
            memcpy(dst,src,xsize);
            src += wrap;
            dst += xsize;
        }
    }
    else if (*pGotPic)	//just out yuv pointer
    {
        for(i=0;i<3;i++)
            {
                frame->yuv[i] = pH264DecTemp->picture->data[i];
            }

        frame->width = pH264DecTemp->c->width;
        frame->height = pH264DecTemp->c->height;
        frame->key_frame = pH264DecTemp->picture->key_frame;	//get key frame flag
        frame->key_frame = pH264DecTemp->picture->key_frame;	//get key frame flag
    }

    return ret;

}
